#include "image_infer_base.h"
#include "vortex/utils/fileops.h"

namespace vortex
{
    ImageInferBase::~ImageInferBase()
    {
        delete m_Context;
        delete m_Engine;
        delete m_Runtime;
    }

    bool ImageInferBase::SetIOInfo(
        const std::vector<BlobInfo>& input_info, 
        const std::vector<BlobInfo>& output_info)
    {
        assert(input_info.size() > 0 && "invalid input info");
        auto shape = input_info[0].shape;
        m_BatchSize = shape[0];

        // create blobs
        for (size_t i = 0; i < input_info.size(); ++i)
        {
            auto name = input_info[i].name;
            auto shape = input_info[i].shape;
            m_InputNames.push_back(name);
            std::shared_ptr<BlobF> blob = std::make_shared<BlobF>(shape);
            m_InputBlobs.push_back(blob);
        }

        for (size_t i = 0; i < output_info.size(); ++i)
        {
            auto name = output_info[i].name;
            auto shape = output_info[i].shape;
            m_OutputNames.push_back(name);
            auto blob = std::make_shared<BlobF>(shape);
            m_OutputBlobs.push_back(blob);
        }
        return true;
    } 

    void ImageInferBase::InternalInfer()
    {
        size_t n_inputs = m_InputBlobs.size();
        size_t n_outputs = m_OutputBlobs.size();
        size_t n_ios = n_inputs + n_outputs;
        std::vector<float*> array_buffer;
        array_buffer.resize(n_ios);

        // feed inputs
        for (size_t i = 0; i < n_inputs; ++i)
        {
            const int input_index = m_Engine->getBindingIndex(m_InputNames[i].c_str());
            if (input_index >= n_ios) continue;
            array_buffer[input_index] = m_InputBlobs[i]->data_device;

            // set binding dimension
            nvinfer1::Dims dims;
            
            // assume the model inputs are 4 dims
            auto shape = m_InputBlobs[i]->shape;
            dims.d[0] = m_BatchSize;
            dims.d[1] = shape[1];
            dims.d[2] = shape[2],
            dims.d[3] = shape[3];
            dims.nbDims = 4;
            m_Context->setBindingDimensions(input_index, dims);
        }

        for (size_t i = 0; i < n_outputs; ++i)
        {
            const int output_index = m_Engine->getBindingIndex(m_OutputNames[i].c_str());
            if (output_index >= n_ios) continue;
            array_buffer[output_index] = m_OutputBlobs[i]->data_device;
        }

        // collect output
        m_Context->enqueue(m_BatchSize, (void**)array_buffer.data(), m_Stream, nullptr);
        cudaStreamSynchronize(m_Stream);
    }

    bool ImageInferBase::LoadEngine(const std::string& engine_path)
    {
        // load data from file
        std::vector<unsigned char> engine_data;
        bool ret = loadBinaryContent(engine_path, engine_data);
        if (!ret)
            return false;

        // create engine
        m_Runtime = nvinfer1::createInferRuntime(m_Logger);
        if (m_Runtime == nullptr)
            return false;
        m_Engine = m_Runtime->deserializeCudaEngine(engine_data.data(), engine_data.size());
        if (m_Engine == nullptr)
            return false;
        m_Context = m_Engine->createExecutionContext();
        if (m_Context == nullptr)
            return false;
        
        // create cuda stream
        checkRuntime(cudaStreamCreate(&m_Stream));
        return true;
    }

}
